from PreRec.BPRMF.InteractScore import InteractScore
import numpy as np
import torch
import torch.utils.data as Data
from PreRec.BPRMF.FileLoader import FileLoader
from PreRec.BPRMF.InteractScore import InteractScore
from tqdm import tqdm
from utils.log_output import simple_text_log


class DataLoader:
    def __init__(self, hyper_params):
        self.hyper_params = hyper_params

        # Load from file
        self.file_loader = FileLoader(hyper_params)
        self.interact_score = InteractScore(hyper_params)

        self.neg_sample_cnt = hyper_params['neg_sample_cnt']
        self.epsilon = hyper_params['epsilon']

    def generate_dataset(self, save_dataset=True) -> Data.Dataset:
        print('BPRMF: Generating Dataset...')
        simple_text_log('train', 'Generating BPR-MF dataset...')
        positive_item = []
        negative_item = []
        user_id = []

        for uid, interact_hist in tqdm(self.file_loader.intereaction_list.items()):
            # Get Positive Negative Pairs
            score_dict = self.interact_score.get_interact_score(
                interact_hist)

            # BPR Dataset
            book_set = set(score_dict.keys())
            not_book_set = set(range(self.file_loader.book_cnt)) - book_set

            book_list = []
            for book in book_set:
                if score_dict[book] > self.epsilon:
                    book_list.append(book)

            not_book_list = list(not_book_set)

            phase2_positive = np.random.choice(
                book_list, size=self.neg_sample_cnt)
            phase2_negative = np.random.choice(
                not_book_list, size=self.neg_sample_cnt)

            positive_item.extend(phase2_positive)
            negative_item.extend(phase2_negative)
            user_id.extend([uid] * self.neg_sample_cnt)

        positive_tensor = torch.LongTensor(positive_item)
        negative_tensor = torch.LongTensor(negative_item)
        user_id_tensor = torch.LongTensor(user_id)

        dataset = Data.TensorDataset(
            positive_tensor, negative_tensor, user_id_tensor)

        print('Finished!')

        return dataset
